/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
package com.webank.tikv.region;

import com.webank.tikv.PDClient;
import com.webank.tikv.TiConfiguration;
import com.webank.tikv.kvproto.Metapb;
import com.webank.tikv.util.ChannelFactory;
import com.webank.tikv.util.ConcreteBackOffer;
import io.grpc.ManagedChannel;
import io.grpc.health.v1.HealthCheckRequest;
import io.grpc.health.v1.HealthCheckResponse;
import io.grpc.health.v1.HealthGrpc;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

public class StoreHealthChecker implements Runnable {

    private static final Logger LOGGER = LogManager.getLogger(StoreHealthChecker.class);

    private static final long MAX_CHECK_STORE_TOMBSTONE_TICK = 60;
    private final BlockingQueue<TiStore> taskQueue;
    private final ChannelFactory channelFactory;
    private final PDClient pdClient;
    private long checkTombstoneTick;
    private final long timeout;

    public StoreHealthChecker(ChannelFactory channelFactory, PDClient pdClient, TiConfiguration conf) {
        this.taskQueue = new LinkedBlockingQueue<>();
        this.channelFactory = channelFactory;
        this.pdClient = pdClient;
        this.checkTombstoneTick = 0;
        this.timeout = conf.getGrpcHealthCheckTimeout();
    }

    // 在刷新到新store时就加入队列, 并定时检测store状态
    // 保存所有非Tombstone的store信息
    public void scheduleStoreHealthCheck(TiStore store) {
        taskQueue.add(store);
    }

    private List<TiStore> getValidStores() {
        List<TiStore> validStores = new ArrayList<>();
        TiStore store;
        while ((store = taskQueue.poll()) != null) {
            try {
                if (!store.isValid()) {
                    continue;
                }
                validStores.add(store);
            } catch (Exception ignore) {
            }
        }
        return validStores;
    }

    private boolean checkStoreHealth(TiStore store) {
        String addressStr = store.getStore().getAddress();
        try {
            ManagedChannel channel = channelFactory.getChannel(addressStr);
            HealthGrpc.HealthBlockingStub stub =
                HealthGrpc.newBlockingStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
            HealthCheckRequest req = HealthCheckRequest.newBuilder().build();
            HealthCheckResponse resp = stub.check(req);
            return resp.getStatus() == HealthCheckResponse.ServingStatus.SERVING;
        } catch (Exception e) {
            return false;
        }
    }

    private boolean checkStoreTombstone(TiStore store) {
        try {
            Metapb.Store newStore = pdClient.getStore(
                ConcreteBackOffer.newCustomBackOff(5000), store.getId());
            if (newStore != null && newStore.getState() == Metapb.StoreState.Tombstone) {
                return true;
            }
        } catch (Exception e) {
            LOGGER.warn("fail to check tombstone stores, ex:{}", e.getMessage());
        }
        return false;
    }

    @Override
    public void run() {
        checkTombstoneTick += 1;
        boolean needCheckTombstoneStore = false;
        if (checkTombstoneTick >= MAX_CHECK_STORE_TOMBSTONE_TICK) {
            needCheckTombstoneStore = true;
            checkTombstoneTick = 0;
        }

        List<TiStore> unReachableStore = new ArrayList<>();
        for (TiStore store : getValidStores()) {
            // tombstone状态的store不再检查
            if (needCheckTombstoneStore && checkStoreTombstone(store)) {
                continue;
            }

            boolean isHealthNow = checkStoreHealth(store);
            if (isHealthNow && !store.isReachable()) {
                // store之前不可达, 但现在已恢复
                LOGGER.warn("store:{} recovers to be reachable", store.getAddress());
                store.markReachable();
            } else if (!isHealthNow && store.isReachable()) {
                // store之前可达, 但现在不可达
                // 先加入unReachableStore, 后续再进行一次检查
                unReachableStore.add(store);
                continue;
            }
            taskQueue.add(store);
        }

        if (unReachableStore.isEmpty()) {
            return;
        }
        try {
            Thread.sleep(timeout);
        } catch (Exception e) {
            taskQueue.addAll(unReachableStore);
            return;
        }

        // 暂停一段时间后再次确认, 若store不可达则重新加入队列
        for (TiStore store : unReachableStore) {
            if (!checkStoreHealth(store)) {
                LOGGER.warn("recheck store:{} is not reachable", store.getAddress());
                store.markUnreachable();
            }
            taskQueue.add(store);
        }
    }
}
